#include "opencv2/opencv.hpp"
#include "cv_bridge/cv_bridge.h"
#include "sensor_msgs/msg/image.hpp"
#include "aim_interfaces/msg/aim_info.hpp"
#include "rclcpp/rclcpp.hpp"
#include <opencv2/dnn.hpp>
#include <vector>
using namespace std;
using namespace cv;

//已知参数
//1、相机内参
static const Mat K = (Mat_<float>(3,3) << 1462.3697,0,398.59394,0,1469.68385,110.68997,0,0,1);
//2、畸变矩阵
static const Mat D = (Mat_<float>(5,1) << 0.003518, -0.311778,-0.016581,0.023682,0.0000);
//3、平移向量
static const Mat T_cam2robot = (Mat_<float>(3,1) << 0.08f,0.0f,0.05f);
//4、旋转向量
static const Mat R_cam2robot = []{
    Mat R;
    Rodrigues(Vec3f(0,60,20)*float(CV_PI/180.0f),R);
    return R;
}();
//5、灯条尺寸
static const float height = 0.08f;
static const float width = 0.16f;
//6、世界系下装甲板坐标(装甲板中心为坐标1)
static const vector<Point3f> P1 = {{-width/2,height/2,0},{width/2,height/2,0},{width/2,-height/2,0},{-width/2,-height/2,0}};
//7、yolo输入尺寸
auto inputSize = Size(640,640);
//8、识别类别
vector<string> classNames = {
    "Default", "Class_1", "Class_2", "Class_3", "Class_4", "Class_5", 
    "Class_6", "Class_7", "Class_8", "Class_9", "Class_10", 
    "Class_11", "Class_12", "Class_13"
};
//9、识别文件路径
string model_path = "/home/vboxuser/armor/src/armor_topic/src/best.onnx";

//一、装甲板坐标获取代码
//1、获取装甲板灯条轮廓
vector<RotatedRect> detect_lightc(const Mat& img_in){
    //声明hsv图像、红/蓝滤波、输出图像
    Mat hsv_img,mask_red1,mask_red2,mask_blue,out_img;
    //转换为hsv图像
    cvtColor(img_in,hsv_img,COLOR_BGR2HSV);
    //定义红蓝滤波
    inRange(hsv_img,Scalar(0,120,70),Scalar(10,255,255),mask_red1);
    inRange(hsv_img,Scalar(156,120,70),Scalar(180,255,255),mask_red2);
    inRange(hsv_img,Scalar(100,120,70),Scalar(124,255,255),mask_blue);
    //定义输出图像
    out_img = mask_red1 | mask_red2 | mask_blue;
    
    //闭运算腐蚀后膨胀
    morphologyEx(out_img, out_img, MORPH_CLOSE,getStructuringElement(MORPH_RECT, {3, 3}), Point(-1,-1), 3);

    //开运算膨胀后腐蚀
    morphologyEx(out_img, out_img, MORPH_OPEN,getStructuringElement(MORPH_RECT, {5, 5}), Point(-1,-1), 1);

    //新建轮廓容器
    vector<vector<Point>> contours;
    findContours(out_img,contours,RETR_EXTERNAL,CHAIN_APPROX_SIMPLE);
    //灯条的一个集合
    vector<RotatedRect> bars;
    //遍历所查找的轮廓
    for(const auto& c : contours){
        //矩形拟合
        RotatedRect r = minAreaRect(c);
        bars.push_back(r);
    }

    // 便于之后操作，将bars从图形的左到右排序
    sort(bars.begin(),bars.end(),[](const RotatedRect& a,const RotatedRect& b){
        return a.center.x < b.center.x;
    });

    return bars;
}
//2、获取装甲板轮廓坐标
vector<Point2f> findArmorPoint(const vector<RotatedRect>& bars){
    //通过暴力匹配的方式进行配对
    for(size_t i = 0;i < bars.size();++i){
        for(size_t j = i + 1;j < bars.size();++j){
            const auto& r1 = bars[i];
            const auto& r2 = bars[j];
            float len1 = max(r1.size.height, r1.size.width);
            float len2 = max(r2.size.height, r2.size.width);
            if (fabs(len1 - len2) / max(len1, len2) > 0.6f){
                continue;
            }
            float d = norm(r1.center-r2.center);
            if(d > (len1+len2)*1.8f){
                continue;
            }
            //角度筛选
            Vec2f v1 = r1.size.width > r1.size.height ? Vec2f(cos(r1.angle*CV_PI/180), sin(r1.angle*CV_PI/180)) : Vec2f(sin(r1.angle*CV_PI/180), -cos(r1.angle*CV_PI/180));
            Vec2f v2 = r2.size.width > r2.size.height ? Vec2f(cos(r2.angle*CV_PI/180), sin(r2.angle*CV_PI/180)) : Vec2f(sin(r2.angle*CV_PI/180), -cos(r2.angle*CV_PI/180));
            if(v1.dot(v2) < 0.85f){
                continue;
            }
            
            //生成四边形角点
            auto vec = [](const RotatedRect& r) -> Point2f{
                Point2f p[4];
                r.points(p);
                Point2f v = (r.size.width > r.size.height) ? (p[1]-p[0]) : (p[2]-p[1]);
                return v*0.5f;
            };
            
            std::vector<Point2f> corners = {
                r1.center + vec(r1),
                r2.center + vec(r2),
                r2.center - vec(r2),
                r1.center - vec(r1)
            };
            return corners; 
        }
    }
    return {};
};

//定义节点
class ArmorShift : public rclcpp::Node{
public:
    ArmorShift(string name) : Node(name){
        
        //加载模型
        if(!initialize(model_path)){
            RCLCPP_ERROR(get_logger(),"模型初始化失败");
        };

        //定义shift订阅者
        shift_subscribe = this->create_subscription<sensor_msgs::msg::Image>("/sensor_img",10,bind(&ArmorShift::shift_callback,this,placeholders::_1));
        //定义shift发布者
        shift_publisher = this->create_publisher<aim_interfaces::msg::AimInfo>("/aim_target",20);
        //定义定时器
        shift_timer = this->create_wall_timer(chrono::milliseconds(1050),bind(&ArmorShift::timer_callback,this));
        shift_timer->cancel(); 
    }
private:
    //声明shift订阅者
    rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr shift_subscribe;
    //声明shift发布者
    rclcpp::Publisher<aim_interfaces::msg::AimInfo>::SharedPtr shift_publisher;
    //声明发布者的定时器
    rclcpp::TimerBase::SharedPtr shift_timer;
    //定义接收图像
    Mat img_in;
    //定义发送结果
    aim_interfaces::msg::AimInfo result;
    //声明dnn模型
    dnn::Net net;
    //定义Detection
    struct Detection {
        int class_id;          // 类别ID
        float confidence;      // 置信度
        Rect box;              // 边界框
        std::string className; // 类别名称
    };
    //定义是否发布
    bool triggered_ = false;

    //定义加载模型函数
    bool initialize(const string& modelPath){
        try{
            // 加载模型
            net = dnn::readNetFromONNX(modelPath);

            // 设置推理后端和目标设备
            net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
            net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);

            return true;
        }catch(const Exception& e){
            RCLCPP_ERROR(this->get_logger(), "加载模型失败 %s", e.what());
            return false;
        }
    }

        /**
    * 图像预处理 - 调整大小并进行归一化
    * 
    * @param image 输入原始图像
    * @return 处理后的Blob数据，适合神经网络输入
    */
    Mat preprocess(const Mat& image){
        // 创建神经网络输入Blob
        Mat blob;

        dnn::blobFromImage(
            image,              // 输入图像
            blob,               // 输出Blob
            1.0/255.0,          // 归一化系数
            inputSize,          // 调整后的尺寸
            cv::Scalar(0,0,0),  // 均值
            true,               // 是否交换RGB通道
            false,              // 是否为单通道图像
            CV_32F              // 输出数据类型
        );

        return blob;
    }

    /**
    * 目标检测主函数 - 执行推理并获取检测结果
    * 
    * @param image 输入图像
    * @param confThreshold 置信度阈值
    * @param nmsThreshold NMS阈值
    * @return 检测结果列表
    */
    vector<Detection> detect(const cv::Mat& image, float confThreshold = 0.25f, float nmsThreshold = 0.45f){
        // 1. 图像预处理
        Mat blob = preprocess(image);

        // 2. 设置网络输入
        net.setInput(blob);

        // 3. 获取输出层名称
        vector<string> outLayerNames = net.getUnconnectedOutLayersNames();

        // 4. 执行前向推理
        vector<Mat> outputs;
        net.forward(outputs, outLayerNames);

        // 5、后处理返回检测结果
        return postprocess(outputs[0], image.size(), confThreshold, nmsThreshold);
    }
    
    /**
    * 后处理函数 - 处理网络输出并转换为检测结果
     * 
     * @param output 网络输出数据
     * @param originalImageSize 原始图像尺寸
     * @param confThreshold 置信度阈值
     * @param nmsThreshold NMS阈值
     * @return 检测结果列表
     */
    vector<Detection> postprocess(const cv::Mat& output, const cv::Size& originalImageSize,float confThreshold, float nmsThreshold){
        vector<Detection> detections;

        // 验证输入数据维度
        if (output.dims < 3){
            RCLCPP_INFO(get_logger(),"输入维度不足，期望3维数据");
            return detections;
        }

        // 计算图像缩放比例，用于将检测框映射回原图
        float xScale = static_cast<float>(originalImageSize.width) / inputSize.width;
        float yScale = static_cast<float>(originalImageSize.height) / inputSize.height;

        vector<int> classIds;
        vector<float> confidences;
        vector<Rect> boxes;

        int num_classes = output.size[1] - 4;  // 类别数量 = 总行数 - 4个坐标值
        int num_anchors = output.size[2];      // 锚点数量（通常是8400）

        // 遍历所有锚点（检测结果）
        for (int col = 0; col < num_anchors; col++){
            // 提取当前锚点的所有类别概率
            vector<float> class_scores(num_classes);
            for (int cls = 0; cls < num_classes; cls++){
                // 索引计算：batch=0, 行=类别索引+4(偏移), 列=当前锚点
                class_scores[cls] = output.ptr<float>(0)[output.size[2] * (cls + 4) + col];
            }
            // 找出最高置信度及其对应的类别
            int classIdPoint = 0;
            float maxScore = 0;
            for (int cls = 0; cls < num_classes; cls++){
                if (class_scores[cls] > maxScore){
                    maxScore = class_scores[cls];
                    classIdPoint = cls;
                }
            }

            // 筛选高于阈值的检测结果
            if (maxScore > confThreshold){
                // 获取边界框坐标 (XYWH格式)
                float x = output.ptr<float>(0)[output.size[2] * 0 + col];  // 中心点x坐标
                float y = output.ptr<float>(0)[output.size[2] * 1 + col];  // 中心点y坐标
                float w = output.ptr<float>(0)[output.size[2] * 2 + col];  // 宽度
                float h = output.ptr<float>(0)[output.size[2] * 3 + col];  // 高度

                // 转换为左上角坐标并缩放到原图大小
                int left = static_cast<int>((x - w/2) * xScale);
                int top = static_cast<int>((y - h/2) * yScale);
                int width = static_cast<int>(w * xScale);
                int height = static_cast<int>(h * yScale);

                // 保存边界框
                boxes.push_back(cv::Rect(left, top, width, height));

                // 验证并保存类别ID
                int classId = classIdPoint;
                if (classId >= 0 && classId < static_cast<int>(classNames.size())){
                    classIds.push_back(classId);
                }else{
                     classIds.push_back(0);  // 使用默认类别
                }

                // 保存置信度
                confidences.push_back(maxScore);
            }
        }
        // 应用非极大值抑制(NMS)，消除重叠框
        std::vector<int> indices;
        if (!boxes.empty()) {
            cv::dnn::NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
        }

        // 整合最终检测结果
        for (size_t i = 0; i < indices.size(); i++) {
            int idx = indices[i];
            Detection detection;
            detection.class_id = classIds[idx];
            detection.confidence = confidences[idx];
            detection.box = boxes[idx];
            detection.className = classNames[detection.class_id];
        
            detections.push_back(detection);
        }

        return detections;
    }

    //定义订阅函数
    void shift_callback(sensor_msgs::msg::Image::SharedPtr msg){
        if(!triggered_){   
            //将通过cv_bridge转换后的信息转换为图像信息
            img_in = cv_bridge::toCvCopy(msg,"bgr8")->image;
            triggered_ = true;
            shift_timer->reset();
        }
    }

    //定义发布函数
    void timer_callback(){
        if (img_in.empty()) {
            RCLCPP_WARN(this->get_logger(), "图像未准备好，跳过本次处理");
            return;
        }

        //获取轮廓
        auto bars = detect_lightc(img_in);
        //获取图像坐标
        auto corners = findArmorPoint(bars);
        if (corners.size() != 4) {
            RCLCPP_WARN(this->get_logger(), "未检测到完整装甲板");
            return;
        }
        //声明相机相关向量
        Mat rvec(3,1,CV_32F);
        Mat tvec(3,1,CV_32F);
        //pnp解算
        solvePnP(P1,corners,K,D,rvec,tvec);
        //将旋转向量转换为旋转矩阵
        Mat R;
        Rodrigues(rvec,R);
        //装甲板年世界系坐标
        Mat P_W = (Mat_<float>(3,1) << 0,0,0);
        //转换为相机下装甲板坐标
        Mat P_C = R * P_W + tvec;
        //转换为机器人下坐标系
        Mat P_R = R_cam2robot * P_C + T_cam2robot;
        //输入结果
        result.coordinate = {
            uint16_t(P_R.at<float>(0)*1000),
            uint16_t(P_R.at<float>(1)*1000),
            uint16_t(P_R.at<float>(2)*1000)
        };
        
        triggered_ = false; 

        // 执行目标检测
        vector<Detection> detections = detect(img_in);
        
        
        if (detections.empty()) {
            RCLCPP_WARN(this->get_logger(), "未检测到目标");
            
        } else {
            result.type = detections[0].class_id;
            shift_publisher ->publish(result);
            RCLCPP_INFO(get_logger(),"消息已经发布");
        }
    }
};

//四、主函数
int main(int argc,char **argv){
    rclcpp::init(argc,argv);
    //定义功能名
    auto node = make_shared<ArmorShift>("armor_shift");
    rclcpp::spin(node);
    rclcpp::shutdown();
}